# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
"""
Implements the Generalized R-CNN framework
"""

import enum
import torch
from torch import nn
import torch.nn.functional as F

from taa_core.structures.image_list import to_image_list

from ..backbone import build_backbone
from ..rpn.rpn import build_rpn
from ..roi_heads.roi_heads import build_roi_heads


class GeneralizedRCNNIA(nn.Module):
    """
    Main class for Generalized R-CNN. Currently supports boxes and masks.
    It consists of three main parts:
    - backbone
    - rpn
    - heads: takes the features + the proposals from the RPN and computes
        detections / masks from it.
    """

    def __init__(self, cfg):
        super(GeneralizedRCNNIA, self).__init__()

        self.backbone = build_backbone(cfg)
        self.rpn = build_rpn(cfg, self.backbone.out_channels)
        self.roi_heads = build_roi_heads(cfg, self.backbone.out_channels)
        self.ia_on = cfg.MODEL.TAA.IA_ON
        self.ia_feat = cfg.MODEL.TAA.IA_FEAT
        self.ia_feat_type = cfg.MODEL.TAA.IA_TYPE
        self.max_iteration = cfg.SOLVER.MAX_ITER
        self.ratio_type = cfg.MODEL.TAA.IA_RATIO_TYPE

    def forward(self, images, targets=None, iteration=None):
        """
        Arguments:
            images (list[Tensor] or ImageList): images to be processed
            targets (list[BoxList]): ground-truth boxes present in the image (optional)

        Returns:
            result (list[BoxList] or dict[Tensor]): the output from the model.
                During training, it returns a dict[Tensor] which contains the losses.
                During testing, it returns list[BoxList] contains additional fields
                like `scores`, `labels` and `mask` (for Mask R-CNN models).

        """
        if self.training and targets is None:
            raise ValueError("In training mode, targets should be passed")
        images = to_image_list(images)
        features = self.backbone(images.tensors)

        if self.ia_on and self.training and iteration is not None:
            ia_features = [f.clone().detach() for f in features]
            max_iteration = self.max_iteration

            if self.ia_feat:
                ratio = self.get_ratio(
                                ratio_type=self.ratio_type, 
                                max_iteration=max_iteration, 
                                iteration=iteration)
                if self.ia_feat_type == 0:
                    assert False
                    masks = self.IA_feat(ia_features, ratio)
                else:
                    masks = self.IA_feat2(ia_features, ratio)
            else:
                if iteration < max_iteration/3*2:
                    masks = self.IA_grad(images, ia_features, targets)
                else:
                    masks = None
        
        if self.ia_on and self.training:
            proposals, proposal_losses = self.rpn(images, features, targets, masks=masks)
        else:
            proposals, proposal_losses = self.rpn(images, features, targets)

        if self.roi_heads:
            x, result, detector_losses = self.roi_heads(features, proposals, targets)
        else:
            # RPN-only models don't have roi_heads
            x = features
            result = proposals
            detector_losses = {}

        if self.training:
            losses = {}
            losses.update(detector_losses)
            losses.update(proposal_losses)
            return losses

        return result

    def get_ratio(self, ratio_type, max_iteration, iteration):
        if ratio_type == 0:
            return 0.5
        if ratio_type == 1:
            if iteration < max_iteration/6 * 1:
                ratio = 0.0
            elif iteration < max_iteration/6 * 2:
                ratio = 0.1
            elif iteration < max_iteration/6 * 3:
                ratio = 0.2
            elif iteration < max_iteration/6 * 4:
                ratio = 0.3
            elif iteration < max_iteration / 6 * 5:
                ratio = 0.4
            else:
                ratio = 0.5
            return ratio
        if ratio_type == 2:
            return float(iteration) / float(max_iteration) / 2.0
        if ratio_type == 3:
            return math.log(1.0 + float(iteration) * 9.0 / float(max_iteration)) / 2.0
        if ratio_type == 4:
            return 1 / (1 + math.exp(-((float(iteration) - float(max_iteration)/2.0))*10.0/float(max_iteration))) / 2.0
        if ratio_type == 5:
            if iteration < max_iteration / 12 * 1:
                ratio = 0.0
            elif iteration < max_iteration / 12 * 2:
                ratio = 0.1
            elif iteration < max_iteration / 12 * 3:
                ratio = 0.2
            elif iteration < max_iteration / 12 * 4:
                ratio = 0.3
            elif iteration < max_iteration / 12 * 5:
                ratio = 0.4
            elif iteration < max_iteration / 12 * 6:
                ratio = 0.5
            elif iteration < max_iteration / 12 * 7:
                ratio = 0.5
            elif iteration < max_iteration / 12 * 8:
                ratio = 0.4
            elif iteration < max_iteration / 12 * 9:
                ratio = 0.3
            elif iteration < max_iteration / 12 * 10:
                ratio = 0.2
            elif iteration < max_iteration / 12 * 11:
                ratio = 0.1
            else:
                ratio = 0.0
            return ratio
        if ratio_type == 6:
            if iteration < max_iteration/2:
                return float(iteration) / float(max_iteration)
            else:
                return float(max_iteration - iteration) / float(max_iteration)
        if ratio_type == 7:
            if iteration < max_iteration/2:
                return math.log(1.0 + float(iteration) * 18.0 / float(max_iteration)) / 2.0
            else:
                return math.log(1.0 + float(max_iteration - iteration) * 18.0 / float(max_iteration)) / 2.0
        if ratio_type == 8:
            if iteration < max_iteration/2:
                return 1 / (1 + math.exp(-((float(iteration) - float(max_iteration)/2.0))*10.0/float(max_iteration))) / 2.0
            else:
                return 1 / (1 + math.exp(-((float(max_iteration - iteration) - float(max_iteration)/2.0))*10.0/float(max_iteration))) / 2.0
        if ratio_type == 9:
            if iteration < 30000:
                ratio = 0.0
            elif iteration < 60000:
                ratio = 0.5
            elif iteration < 90000:
                ratio = 0.0
            elif iteration < 105000:
                ratio = 0.0
            elif iteration < 120000:
                ratio = 0.5
            else:
                ratio = 0.0
            return ratio

    def IA_feat(self, features, ratio):
        ratio = 0.0
        masks_all = []
        for feature in features:
            n, c, h, w = feature.shape

            gap = F.adaptive_avg_pool2d(feature, 1)
            cam = torch.sum(feature * gap, dim=1)

            cam_tmp = cam.view(n, h*w)
            th_idx = int(h*w*ratio)
            th25_mask_value = torch.sort(cam_tmp, dim=1, descending=True)[0][:, th_idx]
            th25_mask_value = th25_mask_value.view(n, 1).expand(n, h*w)
            th25_mask_value = th25_mask_value.reshape(n, h, w)
            mask = torch.where(cam>th25_mask_value, torch.zeros_like(cam), torch.ones_like(cam))
            mask = mask.reshape(n, 1, h, w)
            masks_all.append(mask)
        return masks_all

    def IA_feat2(self, features, ratio):
        masks_all = []
        for feature in features:
            n, c, h, w = feature.shape

            gap = F.adaptive_avg_pool2d(feature, 1)
            gap_c = gap.view(n, c)
            th_idx = int(c * ratio)
            th_channel_mask_value = torch.sort(gap_c, dim=1, descending=True)[0][:, th_idx]
            th_channel_mask_value = th_channel_mask_value.view(n, 1).expand(n, c)
            mask_channel = torch.where(gap_c > th_channel_mask_value,
                            torch.zeros_like(gap_c), 
                            torch.ones_like(gap_c))
            mask_channel = mask_channel.view(n, c, 1, 1).expand(n, c, h, w)

            cam = torch.sum(feature * gap, dim=1)
            cam_tmp = cam.view(n, h*w)
            th_idx = int(h*w*ratio)
            th_spatial_mask_value = torch.sort(cam_tmp, dim=1)[0][:, th_idx]
            th_spatial_mask_value = th_spatial_mask_value.view(n, 1).expand(n, h*w)
            th_spatial_mask_value = th_spatial_mask_value.reshape(n, h, w)
            mask_spatial = torch.where(cam > th_spatial_mask_value, 
                                        torch.zeros_like(cam),
                                        torch.ones_like(cam))
            mask_spatial = mask_spatial.reshape(n, 1, h, w)
            mask_spatial = mask_spatial.expand(n, c, h, w)

            mask = mask_channel * mask_spatial
            masks_all.append(mask)
        
        return masks_all

    def IA_grad(self, images, features, targets):
        with torch.no_grad():
            masks = self.rpn(images, features, targets=targets, ia=True)

        return masks

            



